from gevent import monkey, spawn

monkey.patch_socket()
import requests
from requests.adapters import HTTPAdapter
import re
from RestfulAttack import RestfulAttack

dictionary = dict()
dictionary["driver"] = "WebDriver Hub"
dictionary["tomcat"] = "401 Unauthorized"
dictionary["jenkins"] = "http://eclipse.org/jetty"



class Finger(RestfulAttack):
    def __init__(self, data, port, concurrent):
        super().__init__(data, int(concurrent))
        self.candy = set()
        self.port = str(port)

    # def do_work(self, row):
    #     try:
    #         ip = re.findall(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b", row)[0]
    #         url = "http://" + ip + ":" + self.port + "/wd/hub/static/resource/hub.html"
    #
    #         r = self.get(url=url, timeout=15, allow_redirects=False)
    #
    #         if isinstance(r, Exception):
    #             pass
    #         elif r.status_code != 400 and r.status_code != 302:
    #
    #             server = " ****"
    #
    #             for k, v in dictionary.items():
    #                 if v in r.text:
    #                     server = k
    #                     break
    #             print("***One Service:", server)
    #             self.candy.add(ip + " " + self.port + " " + server + "\n")
    #
    #     except Exception as e:
    #         print(e)
    #     finally:
    #         print(str(self.timer.calls * 100 / len(self._input))[0:4] + "%", " use_time:  ", self.timer.toc())
    #


    def do_work(self, row):
        try:
            ip = re.findall(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b", row)[0]
            url = "http://" + ip + ":" + self.port + "/wd/hub/static/resource/hub.html"

            r = self.get_http1_0(url=url, timeout=8, allow_redirects=False)

            server = " ****"
            for k, v in dictionary.items():
                if v in r:
                    server = k
                    break
            print("***One Service:", server)
            self.candy.add(ip + " " + self.port + " " + server + "\n")

        except Exception as e:
            print("")
        finally:
            now = self.timer.toc()
            print(str(self.timer.calls * 100 / len(self._input))[0:4] + "%",
                  " use_time:  ", now,
                  " total_time:  ", (now/self.timer.calls)*len(self._input))



    def done(self):
        driver = set()
        for i in self.candy:
            if 'driver' in i:
                try:
                    ip = re.findall(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b", i)[0]
                    driver.add(ip + "\n")
                except:
                    pass
        open("driver", 'a+').writelines(driver)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', dest='file', help='port')
    parser.add_argument('-p', dest='port', help='port')
    parser.add_argument('-t', dest='thr', help='thr')


    args = parser.parse_args()

    assert args.file is not None
    assert args.port is not None
    pure_data = open(args.file, 'r').readlines()[1:-1]
    finger = Finger(pure_data, args.port, concurrent=args.thr)
    finger.run()
